package com.distributed.tcp.manager;

import com.distributed.tcp.config.TcpServerProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

@Component
@Slf4j
public class ThreadPoolManager {
    private final Map<Integer, ExecutorService> portThreadPools = new HashMap<>();
    private final TcpServerProperties properties;
    private final ExecutorService threadExecutor;

    @Autowired
    public ThreadPoolManager(TcpServerProperties properties) {
        this.properties = properties;
        // 初始化默认线程池
        TcpServerProperties.ThreadPoolConfig threadPoolConfig = new TcpServerProperties.ThreadPoolConfig();
        this.threadExecutor = createThreadPool(threadPoolConfig, "tcp");
    }

    /**
     * 获取端口对应的线程池
     */
    public ExecutorService getExecutorForPort(int port) {
        // 双重检查锁定确保线程安全
        if (!portThreadPools.containsKey(port)) {
            synchronized (this) {
                if (!portThreadPools.containsKey(port)) {
                    // 获取端口特定配置，若无则使用默认配置
                    TcpServerProperties.ThreadPoolConfig config = 
                        properties.getThreadPools().getOrDefault(port, new TcpServerProperties.ThreadPoolConfig());
                    
                    ExecutorService executor = createThreadPool(config, "port-" + port);
                    portThreadPools.put(port, executor);
                }
            }
        }
        return portThreadPools.get(port);
    }

    /**
     * 创建线程池
     */
    private ExecutorService createThreadPool(TcpServerProperties.ThreadPoolConfig config, String namePrefix) {
        // 参数校验
        if (config.getCoreSize() < 0) {
            throw new IllegalArgumentException("核心线程数不能为负数: " + config.getCoreSize());
        }
        if (config.getMaxSize() <= 0 || config.getMaxSize() < config.getCoreSize()) {
            throw new IllegalArgumentException("最大线程数无效: " + config.getMaxSize());
        }
        if (config.getQueueCapacity() < 0) {
            throw new IllegalArgumentException("队列容量不能为负数: " + config.getQueueCapacity());
        }

        return new ThreadPoolExecutor(
                config.getCoreSize(),
                config.getMaxSize(),
                config.getKeepAliveSeconds(),
                TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(config.getQueueCapacity()),
                new ThreadFactory() {
                    private final AtomicInteger counter = new AtomicInteger(0); // 用AtomicInteger保证线程安全
                    @Override
                    public Thread newThread(Runnable r) {
                        Thread thread = new Thread(r);
                        thread.setName(namePrefix + "-thread-" + counter.getAndIncrement());
                        thread.setDaemon(true);
                        return thread;
                    }
                },
                // 拒绝策略优化：记录更多上下文（如任务内容摘要）
                (r, executor) -> {
                    String taskInfo = r.toString().substring(0, Math.min(r.toString().length(), 100)); // 任务信息摘要
                    log.warn("线程池[{}]任务被拒绝，队列大小: {}, 任务摘要: {}",
                            namePrefix, executor.getQueue().size(), taskInfo);
                    try {
                        if (executor.getQueue().offer(r, 1, TimeUnit.SECONDS)) {
                            log.info("任务在等待后提交成功");
                        } else {
                            log.error("任务提交失败，队列仍满");
                            // 可添加 metrics 监控（如Prometheus记录拒绝次数）
                        }
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                        throw new RejectedExecutionException("任务提交被中断", e);
                    }
                }
        );
    }

    @PreDestroy
    public void shutdownAll() {
        // 关闭端口线程池
        portThreadPools.values().forEach(executor -> {
            executor.shutdown();
            try {
                if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
                    executor.shutdownNow();
                }
            } catch (InterruptedException e) {
                executor.shutdownNow();
            }
        });
        portThreadPools.clear();

        // 关闭默认线程池
        threadExecutor.shutdown();
        try {
            if (!threadExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                threadExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            threadExecutor.shutdownNow();
        }
    }
}
